"""
金融知识图谱构建器
"""

import sqlite3
import networkx as nx
from typing import Dict, List, Any
from utils.logger_config import LoggerManager

logger_manager = LoggerManager(name="FinancialKnowledgeGraph", log_file="financial_kg.log")
logger = logger_manager.logger

class FinancialKnowledgeGraph:
    """金融知识图谱"""
    
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.graph = nx.DiGraph()
        self._initialize_graph()
    
    def _initialize_graph(self):
        """初始化知识图谱"""
        logger.info("正在初始化金融知识图谱...")
        
        try:
            # 连接数据库
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row
            
            # 构建实体
            self._build_entities(conn)
            
            # 构建关系
            self._build_relationships(conn)
            
            conn.close()
            
            logger.info(f"知识图谱构建完成，包含 {len(self.graph.nodes)} 个节点和 {len(self.graph.edges)} 条边")
            
        except Exception as e:
            logger.error(f"初始化知识图谱时出错: {e}")
    
    def _build_entities(self, conn: sqlite3.Connection):
        """构建实体"""
        # 构建基金实体
        self._build_fund_entities(conn)
        
        # 构建股票实体
        self._build_stock_entities(conn)
        
        # 构建债券实体
        self._build_bond_entities(conn)
    
    def _build_fund_entities(self, conn: sqlite3.Connection):
        """构建基金实体"""
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT 基金代码, 基金简称, 基金类型, 管理人, 托管人 FROM 基金基本信息")
            funds = cursor.fetchall()
            
            for fund in funds:
                fund_code = fund['基金代码']
                fund_name = fund['基金简称']
                fund_type = fund['基金类型']
                manager = fund['管理人']
                custodian = fund['托管人']
                
                # 添加基金节点
                self.graph.add_node(
                    f"fund_{fund_code}",
                    type="fund",
                    name=fund_name,
                    code=fund_code,
                    fund_type=fund_type,
                    manager=manager,
                    custodian=custodian
                )
                
        except Exception as e:
            logger.error(f"构建基金实体时出错: {e}")
    
    def _build_stock_entities(self, conn: sqlite3.Connection):
        """构建股票实体"""
        try:
            # 获取A股信息
            cursor = conn.cursor()
            cursor.execute("SELECT DISTINCT 股票代码, 股票名称 FROM A股票日行情表")
            stocks = cursor.fetchall()
            
            for stock in stocks:
                stock_code = stock['股票代码']
                stock_name = stock['股票名称']
                
                # 添加A股节点
                self.graph.add_node(
                    f"stock_{stock_code}",
                    type="stock",
                    name=stock_name,
                    code=stock_code,
                    market="A股"
                )
            
            # 获取港股信息
            cursor.execute("SELECT DISTINCT 股票代码, 股票名称 FROM 港股票日行情表")
            hk_stocks = cursor.fetchall()
            
            for stock in hk_stocks:
                stock_code = stock['股票代码']
                stock_name = stock['股票名称']
                
                # 添加港股节点
                self.graph.add_node(
                    f"hk_stock_{stock_code}",
                    type="stock",
                    name=stock_name,
                    code=stock_code,
                    market="港股"
                )
                
        except Exception as e:
            logger.error(f"构建股票实体时出错: {e}")
    
    def _build_bond_entities(self, conn: sqlite3.Connection):
        """构建债券实体"""
        try:
            # 获取基金债券持仓信息
            cursor = conn.cursor()
            cursor.execute("""
                SELECT DISTINCT 债券代码, 债券名称 
                FROM 基金债券持仓明细 
                WHERE 债券代码 IS NOT NULL AND 债券名称 IS NOT NULL
            """)
            bonds = cursor.fetchall()
            
            for bond in bonds:
                bond_code = bond['债券代码']
                bond_name = bond['债券名称']
                
                # 添加债券节点
                self.graph.add_node(
                    f"bond_{bond_code}",
                    type="bond",
                    name=bond_name,
                    code=bond_code
                )
                
        except Exception as e:
            logger.error(f"构建债券实体时出错: {e}")
    
    def _build_relationships(self, conn: sqlite3.Connection):
        """构建实体间关系"""
        try:
            # 构建基金-股票持仓关系
            self._build_fund_stock_relationships(conn)
            
            # 构建基金-债券持仓关系
            self._build_fund_bond_relationships(conn)
            
            # 构建基金-可转债持仓关系
            self._build_fund_convertible_bond_relationships(conn)
            
        except Exception as e:
            logger.error(f"构建实体关系时出错: {e}")
    
    def _build_fund_stock_relationships(self, conn: sqlite3.Connection):
        """构建基金-股票持仓关系"""
        try:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT 基金代码, 股票代码, 市值, 市值占基金资产净值比 
                FROM 基金股票持仓明细 
                WHERE 基金代码 IS NOT NULL AND 股票代码 IS NOT NULL
            """)
            holdings = cursor.fetchall()
            
            for holding in holdings:
                fund_code = holding['基金代码']
                stock_code = holding['股票代码']
                
                # 检查基金和股票节点是否存在
                fund_node = f"fund_{fund_code}"
                stock_node = f"stock_{stock_code}"
                
                if fund_node in self.graph.nodes and stock_node in self.graph.nodes:
                    # 添加持仓关系，增加字段存在性检查
                    market_value = holding.get('市值', 0)
                    ratio = holding.get('市值占基金资产净值比', 0)
                    
                    self.graph.add_edge(
                        fund_node,
                        stock_node,
                        relation="holds",
                        market_value=market_value,
                        ratio=ratio
                    )
                    
        except Exception as e:
            logger.error(f"构建基金-股票关系时出错: {e}")
    
    def _build_fund_bond_relationships(self, conn: sqlite3.Connection):
        """构建基金-债券持仓关系"""
        try:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT 基金代码, 债券代码, 市值, 市值占基金资产净值比 
                FROM 基金债券持仓明细 
                WHERE 基金代码 IS NOT NULL AND 债券代码 IS NOT NULL
            """)
            holdings = cursor.fetchall()
            
            for holding in holdings:
                fund_code = holding['基金代码']
                bond_code = holding['债券代码']
                
                # 检查基金和债券节点是否存在
                fund_node = f"fund_{fund_code}"
                bond_node = f"bond_{bond_code}"
                
                if fund_node in self.graph.nodes and bond_node in self.graph.nodes:
                    # 添加持仓关系，增加字段存在性检查
                    market_value = holding.get('市值')
                    ratio = holding.get('市值占基金资产净值比')
                    
                    self.graph.add_edge(
                        fund_node,
                        bond_node,
                        relation="holds",
                        market_value=market_value if market_value is not None else 0,
                        ratio=ratio if ratio is not None else 0
                    )
                    
        except Exception as e:
            logger.error(f"构建基金-债券关系时出错: {e}")
    
    def _build_fund_convertible_bond_relationships(self, conn: sqlite3.Connection):
        """构建基金-可转债持仓关系"""
        try:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT 基金代码, 债券代码, 市值, 市值占基金资产净值比 
                FROM 基金可转债持仓明细 
                WHERE 基金代码 IS NOT NULL AND 债券代码 IS NOT NULL
            """)
            holdings = cursor.fetchall()
            
            for holding in holdings:
                fund_code = holding['基金代码']
                bond_code = holding['债券代码']
                
                # 检查基金和债券节点是否存在
                fund_node = f"fund_{fund_code}"
                bond_node = f"bond_{bond_code}"
                
                if fund_node in self.graph.nodes and bond_node in self.graph.nodes:
                    # 添加持仓关系，增加字段存在性检查
                    market_value = holding.get('市值')
                    ratio = holding.get('市值占基金资产净值比')
                    
                    self.graph.add_edge(
                        fund_node,
                        bond_node,
                        relation="holds_convertible_bond",
                        market_value=market_value if market_value is not None else 0,
                        ratio=ratio if ratio is not None else 0
                    )
                    
        except Exception as e:
            logger.error(f"构建基金-可转债关系时出错: {e}")
    
    def query_entity_relationships(self, entity_id: str, max_depth: int = 2) -> Dict:
        """
        查询实体的关系信息
        
        Args:
            entity_id: 实体ID
            max_depth: 最大搜索深度
            
        Returns:
            关系信息字典
        """
        if entity_id not in self.graph.nodes:
            return {"error": f"实体 {entity_id} 不存在"}
        
        # 获取实体基本信息
        entity_info = dict(self.graph.nodes[entity_id])
        
        # 获取直接关联的实体
        neighbors = list(self.graph.neighbors(entity_id))
        relationships = []
        
        for neighbor in neighbors:
            edge_data = self.graph.get_edge_data(entity_id, neighbor)
            neighbor_info = dict(self.graph.nodes[neighbor])
            
            relationships.append({
                "target": neighbor_info,
                "relationship": edge_data
            })
        
        return {
            "entity": entity_info,
            "relationships": relationships
        }
    
    def find_paths_between_entities(self, source_id: str, target_id: str, max_paths: int = 5) -> List[List]:
        """
        查找两个实体之间的路径
        
        Args:
            source_id: 源实体ID
            target_id: 目标实体ID
            max_paths: 最大路径数
            
        Returns:
            路径列表
        """
        try:
            # 使用networkx查找所有简单路径
            paths = list(nx.all_simple_paths(self.graph, source_id, target_id, cutoff=max_paths))
            return paths
        except nx.NetworkXNoPath:
            return []
        except Exception as e:
            logger.error(f"查找路径时出错: {e}")
            return []
    
    def get_subgraph(self, entities: List[str], depth: int = 1) -> nx.DiGraph:
        """
        获取指定实体的子图
        
        Args:
            entities: 实体列表
            depth: 遍历深度
            
        Returns:
            子图
        """
        try:
            # 收集所有相关节点
            related_nodes = set(entities)
            
            for _ in range(depth):
                new_nodes = set()
                for node in related_nodes:
                    # 添加邻居节点
                    neighbors = list(self.graph.neighbors(node)) + list(self.graph.predecessors(node))
                    new_nodes.update(neighbors)
                related_nodes.update(new_nodes)
            
            # 构建子图
            subgraph = self.graph.subgraph(related_nodes).copy()
            return subgraph
        except Exception as e:
            logger.error(f"获取子图时出错: {e}")
            return nx.DiGraph()

from langgraph.graph import StateGraph, END
from app.graph_rag.nodes import GraphRAGNodes
from app.graph_rag.state import GraphState

class GraphRAGBuilder:
    """GraphRAG构建器"""
    
    def __init__(self):
        self.nodes = GraphRAGNodes()
    
    def build_graph(self) -> StateGraph:
        """
        构建GraphRAG图
        
        Returns:
            构建的StateGraph
        """
        # 创建图
        workflow = StateGraph(GraphState)
        
        # 添加节点
        workflow.add_node("determine_processing_type", self.nodes.determine_processing_type)
        workflow.add_node("retrieve_documents", self.nodes.retrieve_documents)
        workflow.add_node("query_finance_database", self.nodes.query_finance_database)
        workflow.add_node("generate_response", self.nodes.generate_response)
        
        # 添加边
        workflow.set_entry_point("determine_processing_type")
        
        # 根据处理类型路由
        workflow.add_conditional_edges(
            "determine_processing_type",
            self._route_based_on_processing_type,
            {
                "finance_agent": "query_finance_database",
                "complex_finance_query": "query_finance_database",
                "prospectus_rag": "retrieve_documents",
                "general_rag": "retrieve_documents",
            }
        )
        
        # 添加文档检索后的边
        workflow.add_edge("retrieve_documents", "generate_response")
        
        # 添加金融数据库查询后的边
        workflow.add_conditional_edges(
            "query_finance_database",
            self._route_after_finance_query,
            {
                "generate_response": "generate_response",
                "retrieve_documents": "retrieve_documents"
            }
        )
        
        # 设置结束点
        workflow.add_edge("generate_response", END)
        
        return workflow
    
    def _route_based_on_processing_type(self, state: GraphState) -> str:
        """
        根据处理类型路由
        
        Args:
            state: 图状态
            
        Returns:
            下一个节点名称
        """
        processing_type = state.get("processing_type", "general_rag")
        return processing_type
    
    def _route_after_finance_query(self, state: GraphState) -> str:
        """
        金融查询后的路由
        
        Args:
            state: 图状态
            
        Returns:
            下一个节点名称
        """
        query_result = state.get("query_result", "")
        processing_type = state.get("processing_type", "general_rag")
        intent_data = state.get("intent_data", {})
        
        # 如果查询结果为空，且意图是company_background，则路由到retrieve_documents
        if (not query_result and 
            intent_data.get("intent") == "company_background"):
            # 更新处理类型确保retrieve_documents能正确识别
            state["processing_type"] = "prospectus_rag"
            logger.info("检测到company_background意图，更新处理类型为prospectus_rag")
            return "retrieve_documents"
        # 如果查询结果为空，但不是company_background意图，则也路由到retrieve_documents
        elif not query_result:
            return "retrieve_documents"
        else:
            return "generate_response"

# 创建全局实例
graph_builder = GraphRAGBuilder()
